"""
 微调澜舟科技开源的mengzi-t5-base做摘要生成
    数据集(只取5000条)：nlpcc_2017: https://huggingface.co/datasets/supremezxc/nlpcc_2017
    model link：https://modelscope.cn/models/langboat/mengzi-t5-base
          link: https://huggingface.co/Langboat/mengzi-t5-base

    孟子中文T5预训练生成模型与T5结构相同，只有无监督数据训练，不包含下游任务，需要在特定任务上finetune后使用
"""
import os
import platform
import torch
import torch.nn as nn
from torch.utils.data import Dataset, DataLoader
from datasets import load_from_disk
from transformers import AdamW
from transformers import T5Tokenizer, T5ForConditionalGeneration
# from logging_util import get_logger
from rouge import Rouge


device = 'cuda' if torch.cuda.is_available() else 'cpu'

# 获取当前操作系统的名称
os_name = platform.system()
# logger = get_logger(model_name='mengzi-t5-base')

# 设置模型路径及数据集路径
if os_name == "Windows":
    model_dir = r'D:\python\models\langboat\meng_zi_t5'
    data_dir = r'D:\python\datas\nlp_seq2seq\nlpcc_2017'
    logger.info("当前执行环境是 Windows...")
elif os_name == "Linux":
    model_dir = r'/root/autodl-fs/models/meng_zi_t5'
    data_dir = r'/root/autodl-fs/data/nlp_ai/nlp_seq2seq/nlpcc_2017'
    logger.info("当前执行环境是 Linux...")
else:
    raise ValueError("当前执行环境不是 Windows 也不是 Linux")


class Dataset(Dataset):
    def __init__(self, split, num=None):

        # 在线加载数据集(需外网环境)
        # dataset = load_dataset(path='nlpcc_2017', split=split)
        # 我们可以将数据保存到本地磁盘，下次利用load_from_disk直接从本地加载即可
        # dataset.save_to_disk("./nlpcc_2017")

        # 划分数据集为：训练集和测试集
        dataset = load_from_disk(dataset_path=data_dir)
        # 选取4900条训练集、100条测试集
        split_dataset = dataset.train_test_split(100, seed=42)
        if num:
            # 离线加载数据集，cpu环境取少量数据用来训练、测试模型
            dataset = split_dataset[split].select(range(num))
        else:
            # 离线加载数据集
            dataset = split_dataset[split]

        # 过滤掉太长的句子，需要去掉CLS、SEP
        def f(data):
            return len(data['content']) <= 512 - 2

        dataset = dataset.filter(f)

        self.dataset = dataset

    def __len__(self):
        return len(self.dataset)

    def __getitem__(self, i):
        content = self.dataset[i]['content']
        title = self.dataset[i]['title']
        return content, title

    # 对数据集进行组装
    # 主要就是在content前面加上特定的提示词
    # 然后调用tokenizer进行批处理
    def get_collate_fn(tokenizer):
        def collate_fn(batch):
            contents = ["摘要生成: \n" + tup2[0] for tup2 in batch]
            original_labels = [tup2[1] for tup2 in batch]
            # 特殊字符
            # 0 -> <pad>
            # 1 -> </s>
            # 2 -> <unk>
            inputs = tokenizer(contents, max_length=384, truncation=True, return_tensors='pt', padding=True)
            labels = tokenizer(text_target=original_labels, max_length=64, truncation=True, return_tensors='pt',
                               padding=True)
            return inputs, labels

        return collate_fn

    
class MengZiT5Model(nn.Module):
    def __init__(self):
        super().__init__()
        # 加载预训练模型
        self.model = T5ForConditionalGeneration.from_pretrained(model_dir)


    def forward(self, inputs, labels=None):
        # 1、encoder的input_ids和attention_mask
        input_ids = inputs['input_ids']
        attention_mask = inputs['attention_mask']

        if labels is not None:
            # 2、decoder 的labels
            train_labels = labels['input_ids'].contiguous()
            train_labels_mask = labels['attention_mask']

            # 3、decoder 的input_ids和attention_mask
            decoder_input_ids = train_labels.new_zeros(train_labels.shape)
            decoder_input_ids[..., 1:] = train_labels[..., :-1].clone()

            decoder_attention_mask = train_labels_mask.new_zeros(train_labels_mask.shape)
            decoder_attention_mask[..., 1:] = train_labels_mask[..., :-1].clone()
            decoder_attention_mask[..., 0] = 1
            # 4、送入模型进行预测
            outputs = self.model(input_ids=input_ids
                                 , attention_mask=attention_mask
                                 , decoder_input_ids=decoder_input_ids
                                 , decoder_attention_mask=decoder_attention_mask
                                 , labels=train_labels)
            # 5、返回训练时候的Loss值
            return outputs.loss
        else:
            # 模型生成
            summary_ids = self.model.generate(input_ids
                                              , num_beams=4 # 束搜索法
                                              , no_repeat_ngram_size=2 # 确保不重复
                                              , min_length=10 # 长度限制
                                              , max_length=64
                                              , early_stopping=True
            )
            # 将id转换为输出 summary_ids.shape = [bs, length]
            outputs = tokenizer.batch_decode(summary_ids, skip_special_tokens=True)
            return outputs
# 模型训练

